'''
    该文件用于将 xmlFile.py生成的文件整个xml文件 转成 json文件
    json文件格式参考作者工程中的 json文件格式

    转换后person_id的定义为，按人数从1一直到总人数
    iscrowd恒定为0
    segmentation恒为[]
'''
import xml.etree.ElementTree as ET
import json
import copy
import math
# 使用前需要在xml文件前加<root>和</root>
def xml_to_json(xml_path,json_path,img_format='.jpg'):
    # json 文件主要两项内容
    json_dict = dict()
    annotations = list()
    images = list()
    categories = list()
    # 一个标签和一张图
    one_annotation = dict()
    one_image = dict()


    tree = ET.parse(xml_path)
    root = tree.getroot()
    # 获取xml文件相应内容
    for ann in root:
        person_index = 0
        for child in ann:
            if (child.tag == 'filename'):
                one_image['file_name'] = child.text+img_format
                one_image['id'] = child.text
                one_annotation['image_id'] = child.text
            elif (child.tag == 'size'):
                for item in child:
                    if item.tag == 'width':
                        one_image['width'] = int(item.text)
                    elif item.tag == 'height':
                        one_image['height'] = int(item.text)
            elif (child.tag == 'object'):      #  检测对象
                person_index = person_index+1
                one_person_box = list()
                for item in child:   # type,name,pose,truncated,difficult,robndbox.
                    if item.tag == 'name':
                        one_annotation['category_id'] = int(item.text)
                        one_annotation['person_id'] = person_index
                        one_annotation['iscrowd'] = 0             # 恒为0
                        one_annotation['segmentation'] = []       # 恒为[]
                    elif item.tag == 'robndbox':
                        for box in item:
                            one_person_box.append(float(box.text))
                        # 对角度进行变换，将弧度转换成角度，同时和RAPiD论文中的角度对应的情况
                        # 如果rolabelImg标注数据h>w，则角度值小于90度的不变换，大于等于90度的结果减去180度
                        # 如果rolabelImg标注数据w>h，则角度值小于90度的减去90°，大于等于90度的结果减去90°
                        if one_person_box[2]>one_person_box[3]:        # w>h,同时需要交换w和h的值
                            temp = one_person_box[2]
                            one_person_box[2] = one_person_box[3]
                            one_person_box[3] = temp
                            one_person_box[4] = one_person_box[4]*180/math.pi-90
                        else:                 # h>w
                            angle = one_person_box[4]*180/math.pi
                            one_person_box[4] = angle if (angle<90) else (angle-180)
                        # 由于浮点计算存在误差，需要对数据进行处理
                        if (one_person_box[4]>=90) or (one_person_box[4] < -90):   # 大于等于90度，或者小于-90度 情况
                            one_person_box[4] = -90
                one_annotation['bbox'] = one_person_box
                #print(one_annotation['image_id'])
                one_annotation['area'] = one_person_box[2]*one_person_box[3]   # w*,h
                annotations.append(copy.deepcopy(one_annotation))
        images.append(copy.deepcopy(one_image))
    # 类别信息
    category = dict()
    category['supercategory'] ='person'
    category['person'] = 'person'
    category['id'] = 1
    categories.append(category)
    json_dict['annotations'] = annotations
    json_dict['images'] = images
    json_dict['categories']=categories
    # 将获取的xml内容写入到json文件中
    with open(json_path, 'w') as f:
        f.write(json.dumps(json_dict,indent=1,separators=(',',':')))

# json 文件转换成xml文件
from xml.dom.minidom import Document
import os
from tools.get_data_info import get_json_data
def single_json_to_xml(json_path,xml_path,img_format='.jpg'):
    json_info, _ = get_json_data(json_path)
    for file in json_info:
        doc=Document()
        annotation = doc.createElement('annotation')

        doc.appendChild(annotation)

        folder = doc.createElement('folder')
        folder.appendChild(doc.createTextNode('video_mat'))
        annotation.appendChild(folder)

        filename = doc.createElement('filename')
        filename.appendChild(doc.createTextNode(file['filename']))
        annotation.appendChild(filename)

        path = doc.createElement('path')
        path.appendChild(doc.createTextNode(file['filename']+'.jpg'))
        annotation.appendChild(path)

        source = doc.createElement('source')
        database = doc.createElement('database')
        database.appendChild(doc.createTextNode('Unknown'))
        source.appendChild(database)
        annotation.appendChild(source)

        size = doc.createElement('size')
        width = doc.createElement('width')
        width.appendChild(doc.createTextNode(str(file['width'])))
        height = doc.createElement('height')
        height.appendChild(doc.createTextNode(str(file['height'])))
        depth = doc.createElement('depth')
        depth.appendChild(doc.createTextNode('3'))
        size.appendChild(width)
        size.appendChild(height)
        size.appendChild(depth)
        annotation.appendChild(size)

        segmented = doc.createElement('segmented')
        segmented.appendChild(doc.createTextNode('0'))
        annotation.appendChild(segmented)
        for box in file['boxes']:
            object = doc.createElement('object')

            type = doc.createElement('type')
            type.appendChild(doc.createTextNode('robndbox'))

            name = doc.createElement('name')
            name.appendChild(doc.createTextNode('1'))
            pose = doc.createElement('pose')
            pose.appendChild(doc.createTextNode('Unspecified'))
            truncated = doc.createElement('truncated')
            truncated.appendChild(doc.createTextNode('0'))
            difficult = doc.createElement('difficult')
            difficult.appendChild(doc.createTextNode('0'))

            robndbox = doc.createElement('robndbox')
            cx = doc.createElement('cx')
            cx.appendChild(doc.createTextNode(str(box[0])))
            cy = doc.createElement('cy')
            cy.appendChild(doc.createTextNode(str(box[1])))
            w = doc.createElement('w')
            w.appendChild(doc.createTextNode(str(box[2])))
            h = doc.createElement('h')
            h.appendChild(doc.createTextNode(str(box[3])))
            angle= doc.createElement('angle')
            angle.appendChild(doc.createTextNode(str(box[4]*math.pi/180.0)))
            robndbox.appendChild(cx)
            robndbox.appendChild(cy)
            robndbox.appendChild(w)
            robndbox.appendChild(h)
            robndbox.appendChild(angle)

            object.appendChild(type)
            object.appendChild(name)
            object.appendChild(pose)
            object.appendChild(truncated)
            object.appendChild(difficult)
            object.appendChild(robndbox)
            annotation.appendChild(object)

        f_xml = open(os.path.join(xml_path, file['filename'] + '.xml'), 'w')
        #doc.getRootElement().writexml(f_xml, indent='\t', addindent='\t', newl='\n')
        doc.writexml(f_xml, '\n', '\t', '', 'UTF-8')
        f_xml.close()
# json 文件转换成xml文件


import os
def single_json_to_xml2(json_path,xml_path,img_format='.jpg'):
    json_info, _ = get_json_data(json_path)
    for file in json_info:
        annotation = ET.Element('annotation')
        folder = ET.SubElement(annotation,'folder')
        folder.text = 'video_mat'


        filename = ET.SubElement(annotation,'filename')
        filename.text=file['filename']

        path = ET.SubElement(annotation,'path')
        path.text = file['filename'] + '.jpg'

        source = ET.SubElement(annotation,'source')

        database = ET.SubElement(source,'database')
        database.text = 'Unknown'

        size = ET.SubElement(annotation,'size')
        width = ET.SubElement(size,'width')
        width.text = str(file['width'])
        height = ET.SubElement(size, 'height')
        height.text = str(file['height'])
        depth = ET.SubElement(size, 'depth')
        depth.text = '3'


        segmented = ET.SubElement(annotation,'segmented')
        segmented.text = '0'

        for box in file['boxes']:
            object = ET.SubElement(annotation,'object')

            type = ET.SubElement(object,'type')
            type.text = 'robndbox'

            name = ET.SubElement(object, 'name')
            name.text = '1'

            pose = ET.SubElement(object, 'pose')
            pose.text = 'Unspecified'

            truncated = ET.SubElement(object, 'truncated')
            truncated.text = '0'

            difficult = ET.SubElement(object, 'difficult')
            difficult.text = '0'

            robndbox = ET.SubElement(object, 'robndbox')

            cx = ET.SubElement(robndbox, 'cx')
            cx.text = str(box[0])
            cy = ET.SubElement(robndbox, 'cy')
            cy.text = str(box[1])
            w = ET.SubElement(robndbox, 'w')
            w.text = str(box[2])
            h = ET.SubElement(robndbox, 'h')
            h.text = str(box[3])
            angle = ET.SubElement(robndbox, 'angle')
            angle.text = str(box[4])
        tree = ET.ElementTree(annotation)
        tree.write(os.path.join(xml_path, file['filename'] + '.xml'), encoding='utf-8')

if __name__ == '__main__':
    ''''#   所有文件
    filename = 'Market1'
    xml_path = '../COSSY/annotations/'+filename+'.xml'
    json_path = '../COSSY/annotations/'+filename+'.json'
    xml_to_json(xml_path,json_path)
    '''

    '''
     #单个xml文件,使用前需要在xml文件开头加<root>,在xml文件末尾加</root>
    filename = 'Market1_1219'
    xml_path = '../COSSY/test/'+filename+'.xml'
    json_path = '../COSSY/test/'+filename+'.json'
    xml_to_json(xml_path,json_path)
    '''
    # json to xml
    json_path = '../COSSY/test/Market1_1219.json'
    xml_path = '../COSSY/test/'
    single_json_to_xml2(json_path, xml_path)

